package doubao

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	_ "github.com/volcengine/volcengine-go-sdk/volcengine"
)

func init() {
	os.Setenv("ARK_API_KEY", "3f0a2be4-be46-4fd6-8f98-55bce888b637")

}

func client() *arkruntime.Client {
	client := arkruntime.NewClientWithApiKey(
		//通过 os.Getenv 从环境变量中获取 ARK_API_KEY
		os.Getenv("ARK_API_KEY"),
		//深度推理模型耗费时间会较长，请您设置较大的超时时间，避免超时导致任务失败。推荐30分钟以上
		arkruntime.WithTimeout(30*time.Minute),
	)
	return client
}

type JsonReq struct {
	KeyWord string `json:"keyWord"`
}

func Dialog(c *gin.Context) {
	client := client()
	var json JsonReq
	if err := c.ShouldBind(&json); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}


	// 创建一个上下文，通常用于传递请求的上下文信息，如超时、取消等
	ctx := context.Background()
	// 构建聊天完成请求，设置请求的模型和消息内容
	req := model.ChatCompletionRequest{
		// 需要替换 <Model> 为模型的Model ID
		Model: "deepseek-r1-250120",
		Messages: []*model.ChatCompletionMessage{
			{
				// 消息的角色为用户
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String(json.KeyWord),
				},
			},
		},
	}

	stream, err := client.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("stream chat error: %v\n", err)
		return
	}
	defer stream.Close()
	for {
		recv, err := stream.Recv()
		if err == io.EOF {
			return
		}
		if err != nil {
			fmt.Printf("Stream chat error: %v", err)
			return
		}

		if len(recv.Choices) > 0 {
			// 输出 ReasoningContent
			if recv.Choices[0].Delta.ReasoningContent != nil {
				rr := *recv.Choices[0].Delta.ReasoningContent
				fmt.Print(rr)
				c.Writer.Write([]byte(rr))
				c.Writer.Flush()
			}
			// 输出 Delta.Content
			rrr := recv.Choices[0].Delta.Content
			fmt.Print(rrr)
			c.Writer.Write([]byte(rrr))
			c.Writer.Flush()
		}
	}

}

func normal_func() {
	client := arkruntime.NewClientWithApiKey(
		//通过 os.Getenv 从环境变量中获取 ARK_API_KEY
		os.Getenv("ARK_API_KEY"),
		//深度推理模型耗费时间会较长，请您设置较大的超时时间，避免超时导致任务失败。推荐30分钟以上
		arkruntime.WithTimeout(30*time.Minute),
	)
	// 创建一个上下文，通常用于传递请求的上下文信息，如超时、取消等
	ctx := context.Background()
	// 构建聊天完成请求，设置请求的模型和消息内容
	req := model.ChatCompletionRequest{
		// 需要替换 <Model> 为模型的Model ID
		Model: "deepseek-r1-250120",
		Messages: []*model.ChatCompletionMessage{
			{
				// 消息的角色为用户
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("请指出哪吒2电影和大圣归来电影的有哪些优秀之处"),
				},
			},
		},
	}

	// 发送聊天完成请求，并将结果存储在 resp 中，将可能出现的错误存储在 err 中
	resp, err := client.CreateChatCompletion(ctx, req)
	if err != nil {
		// 若出现错误，打印错误信息并终止程序
		fmt.Printf("standard chat error: %v\n", err)
		return
	}
	// 检查是否触发深度推理，触发则打印思维链内容
	if resp.Choices[0].Message.ReasoningContent != nil {
		fmt.Println(*resp.Choices[0].Message.ReasoningContent)
	}
	// 打印聊天完成请求的响应结果
	fmt.Println(*resp.Choices[0].Message.Content.StringValue)
}

func stream_func() {
	client := arkruntime.NewClientWithApiKey(
		//通过 os.Getenv 从环境变量中获取 ARK_API_KEY
		os.Getenv("ARK_API_KEY"),
		//深度推理模型耗费时间会较长，请您设置较大的超时时间，避免超时导致任务失败。推荐30分钟以上
		arkruntime.WithTimeout(30*time.Minute),
	)
	// 创建一个上下文，通常用于传递请求的上下文信息，如超时、取消等
	ctx := context.Background()
	// 构建聊天完成请求，设置请求的模型和消息内容
	req := model.ChatCompletionRequest{
		// 需要替换 <Model> 为模型的Model ID
		Model: "deepseek-r1-250120",
		Messages: []*model.ChatCompletionMessage{
			{
				// 消息的角色为用户
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("2025年全年的放假和法定加班安排时间表"),
				},
			},
		},
	}

	stream, err := client.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("stream chat error: %v\n", err)
		return
	}
	defer stream.Close()

	for {
		recv, err := stream.Recv()
		if err == io.EOF {
			return
		}
		if err != nil {
			fmt.Printf("Stream chat error: %v", err)
			return
		}

		if len(recv.Choices) > 0 {
			// 输出 ReasoningContent
			if recv.Choices[0].Delta.ReasoningContent != nil {
				fmt.Print(*recv.Choices[0].Delta.ReasoningContent)
			}
			// 输出 Delta.Content
			fmt.Print(recv.Choices[0].Delta.Content)
		}
	}
}
